import akshare as ak
import pandas as pd
from sqlalchemy import create_engine
from sqlalchemy import text

# MySQL 连接配置
engine = create_engine('mysql+pymysql://root:123456@localhost:3306/cn_stock_data')

# 生成 2012-2022 年所有季度日期
years = range(2012, 2023)  # 2012 ~ 2022
quarters = ["0331", "0630", "0930", "1231"]
report_dates = [f"{year}{q}" for year in years for q in quarters]

def process_and_save(report_date_str: str):
    print(f"开始处理 {report_date_str} ...")
    try:
        # 1. 下载资产负债表数据
        df = ak.stock_zcfz_em(date=report_date_str)

        if df.empty:
            print(f"⚠️ {report_date_str} 没有数据，跳过")
            return

        # 2. 删除序号列
        if '序号' in df.columns:
            df = df.drop(columns=['序号'])

        # 3. 数字列处理
        numeric_cols = [
            '资产-货币资金', '资产-应收账款', '资产-存货', '资产-总资产',
            '资产-总资产同比', '负债-应付账款', '负债-预收账款',
            '负债-总负债', '负债-总负债同比', '资产负债率', '股东权益合计'
        ]
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')

        # 4. 日期列处理
        if '公告日期' in df.columns:
            df['公告日期'] = pd.to_datetime(df['公告日期'], errors='coerce')

        # 5. 新增 report_date 列
        df['report_date'] = pd.to_datetime(report_date_str, format='%Y%m%d')

        # 6. 重命名列与数据库一致
        df = df.rename(columns={
            '股票代码': 'stock_code',
            '股票简称': 'stock_name',
            '资产-货币资金': 'cash_and_equivalents',
            '资产-应收账款': 'accounts_receivable',
            '资产-存货': 'inventory',
            '资产-总资产': 'total_assets',
            '资产-总资产同比': 'total_assets_yoy',
            '负债-应付账款': 'accounts_payable',
            '负债-预收账款': 'advances_from_customers',
            '负债-总负债': 'total_liabilities',
            '负债-总负债同比': 'total_liabilities_yoy',
            '资产负债率': 'debt_to_assets_ratio',
            '股东权益合计': 'total_shareholders_equity',
            '公告日期': 'announcement_date'
        })

        

        with engine.begin() as conn:
            conn.execute(
                text("DELETE FROM balance_sheet WHERE report_date = :report_date"),
                {"report_date": pd.to_datetime(report_date_str, format="%Y%m%d")}
            )
        df.to_sql('balance_sheet', con=engine, if_exists='append', index=False)

        print(f"✅ {report_date_str} 导入完成！ 共 {len(df)} 行")

    except Exception as e:
        print(f"❌ {report_date_str} 处理失败: {e}")

# 批量处理
for date in report_dates:
    process_and_save(date)

print("🎉 2012-2022 所有季度导入完成！")
